#include <sys/console.h>
#include <asm/processor.h>
#include "asm/msr-index.h"
#include "asm/msr.h"
#include "asm/vmx.h"
#include <asm/cpufeature.h>
#include <asm/desc.h>

#define NR_AUTOLOAD_MSRS        8
#define NR_VCPU_REGS            16


uint32_t revision_id;

char host_stack[PAGE_SIZE] __aligned(PAGE_SIZE);
char guest_stack[PAGE_SIZE] __aligned(PAGE_SIZE);

extern const uint64_t vm_exit_func;
extern const uint64_t guest_entry;

struct vmcs {
    uint32_t revision_id;
    uint32_t abort;
    uint8_t data[PAGE_SIZE - 8];
} __packed __aligned(PAGE_SIZE);

struct kvm_vcpu {
    uint64_t regs[NR_VCPU_REGS];
    uint64_t cr2;
    _Atomic int activity_state;
    uint8_t sipi_vector;
};

struct vcpu_vmx {
    struct kvm_vcpu vcpu;
    uint64_t host_rsp;
    int fail;
    int launched;
    struct msr_autoload {
        struct vmx_msr_entry guest[NR_AUTOLOAD_MSRS];
        struct vmx_msr_entry host[NR_AUTOLOAD_MSRS];
    } msr_autoload;
};

void just_for_break();

static DEFINE_PER_CPU(struct vmcs, vmxarea) __aligned(PAGE_SIZE);
static DEFINE_PER_CPU(struct vmcs, current_vmcs) __aligned(PAGE_SIZE);
static DEFINE_PER_CPU(struct vcpu_vmx, current_vmx);

#define vmcs_check16(field)                                                     \
({                                                                              \
        BUILD_BUG_ON_MSG(((field) & 0x6001) == 0x2000,                          \
                         "16-bit accessor invalid for 64-bit field");           \
        BUILD_BUG_ON_MSG(((field) & 0x6001) == 0x2001,                          \
                         "16-bit accessor invalid for 64-bit high field");      \
        BUILD_BUG_ON_MSG(((field) & 0x6000) == 0x4000,                          \
                         "16-bit accessor invalid for 32-bit high field");      \
        BUILD_BUG_ON_MSG(((field) & 0x6000) == 0x6000,                          \
                         "16-bit accessor invalid for natural width field");    \
})

#define vmcs_check32(field)                                                     \
({                                                                              \
        BUILD_BUG_ON_MSG(((field) & 0x6000) == 0,                               \
                         "32-bit accessor invalid for 16-bit field");           \
        BUILD_BUG_ON_MSG(((field) & 0x6000) == 0x6000,                          \
                         "32-bit accessor invalid for natural width field");    \
})

#define vmcs_check64(field)                                                     \
({                                                                              \
        BUILD_BUG_ON_MSG(((field) & 0x6000) == 0,                               \
                         "64-bit accessor invalid for 16-bit field");           \
        BUILD_BUG_ON_MSG(((field) & 0x6001) == 0x2001,                          \
                         "64-bit accessor invalid for 64-bit high field");      \
        BUILD_BUG_ON_MSG(((field) & 0x6000) == 0x4000,                          \
                         "64-bit accessor invalid for 32-bit field");           \
        BUILD_BUG_ON_MSG(((field) & 0x6000) == 0x6000,                          \
                         "64-bit accessor invalid for natural width field");    \
})

#define vmcs_checkl(field)                                                      \
({                                                                              \
        BUILD_BUG_ON_MSG(((field) & 0x6000) == 0,                               \
                         "Natural width accessor invalid for 16-bit field");    \
        BUILD_BUG_ON_MSG(((field) & 0x6001) == 0x2000,                          \
                         "Natural width accessor invalid for 64-bit field");    \
        BUILD_BUG_ON_MSG(((field) & 0x6001) == 0x2001,                          \
                         "Natural width accessor invalid for 64-bit high field");\
        BUILD_BUG_ON_MSG(((field) & 0x6000) == 0x4000,                          \
                         "Natural width accessor invalid for 32-bit field");    \
})

static __always_inline unsigned long __vmcs_read(unsigned long field) {
    unsigned long value;

    asm volatile("vmread %1, %0"
            : "=a" (value) : "d" (field) : "cc");
    return value;
}

#define vmcs_read16(field)              \
({                                      \
        vmcs_check16(field);            \
        (uint16_t)__vmcs_read(field);   \
})

#define vmcs_read32(field)              \
({                                      \
        vmcs_check32(field);            \
        (uint32_t)__vmcs_read(field);   \
})

#define vmcs_read64(field)              \
({                                      \
        vmcs_check64(field);            \
        (uint64_t)__vmcs_read(field);   \
})

#define vmcs_readl(field)               \
({                                      \
        vmcs_checkl(field);             \
        __vmcs_read(field);             \
})

static __always_inline void __vmcs_write(unsigned long field, unsigned long value) {
    uint8_t error;

    asm volatile("vmwrite %1, %2; setna %0"
            : "=q" (error) : "a" (value), "d" (field) : "cc");
    if (error)
        panic("vmx: vmwrite error: reg %lx value %lx (err %d)\n",
              field, value, vmcs_read32(VM_INSTRUCTION_ERROR));
}

#define vmcs_write16(field, value)      \
({                                      \
        vmcs_check16(field);            \
        __vmcs_write(field, value);     \
})

#define vmcs_write32(field, value)      \
({                                      \
        vmcs_check32(field);            \
        __vmcs_write(field, value);     \
})

#define vmcs_write64(field, value)      \
({                                      \
        vmcs_check64(field);            \
        __vmcs_write(field, value);     \
})

#define vmcs_writel(field, value)       \
({                                      \
        vmcs_checkl(field);             \
        __vmcs_write(field, value);     \
})

uint32_t cpu_has_vmx(void) {
    uint32_t ecx = cpuid_ecx(1);
    return ecx >> 5 & 1;
}

static void native_vmxon(uint64_t addr) {
    asm volatile("vmxon %0" : : "m" (addr) : "memory", "cc");
}

static void vmcs_load(uint64_t addr) {
    uint8_t error;

    asm volatile ("vmptrld %1; setna %0"
            : "=qm" (error) : "m" (addr)
            : "cc", "memory");
    if (error)
        panic("vmx: vmptrld %" PRIx64 "failed\n", addr);
}

static void vmcs_clear(uint64_t addr) {
    uint8_t error;

    asm volatile ("vmclear %1; setna %0"
            : "=qm" (error) : "m" (addr)
            : "cc", "memory");
    if (error)
        panic("vmx: vmptrld %" PRIx64 "failed\n", addr);
}

uint16_t read_es(void) {
    uint16_t val;
    asm volatile("mov %%es, %0\n\t" : "=r"(val));
    return val;
}

uint16_t read_cs(void) {
    uint16_t val;
    asm volatile("mov %%cs, %0\n\t" : "=r"(val));
    return val;
}

uint16_t read_ds(void) {
    uint16_t val;
    asm volatile("mov %%ds, %0\n\t" : "=r"(val));
    return val;
}

uint16_t read_ss(void) {
    uint16_t val;
    asm volatile("mov %%ss, %0\n\t" : "=r"(val));
    return val;
}



static int adjust_vmx_controls(uint32_t ctl_min, uint32_t ctl_opt,
                               uint32_t msr, uint32_t *result) {
    uint32_t vmx_msr_low, vmx_msr_high;
    uint32_t ctl = ctl_min | ctl_opt;

    rdmsr(msr, &vmx_msr_low, &vmx_msr_high);

    ctl &= vmx_msr_high; /* bit == 0 in high word ==> must be zero */
    ctl |= vmx_msr_low;  /* bit == 1 in low word  ==> must be one  */

    /* Ensure minimum (required) set of control bits are supported. */
    if (ctl_min & ~ctl)
        panic("vmx: failed to set minimum control bits\n");

    *result = ctl;
    return 0;
}


int vmx_enable() {
    uint64_t cr4 = read_cr4();
    pr_info("cr4: 0x%lx\n", cr4);

    if (cr4 & X86_CR4_VMXE) {
        panic("vmx: VMXE already set in CR4\n");
    }

    uint64_t old, test_bits;
    old = rdmsrl(MSR_IA32_FEATURE_CONTROL);

    test_bits = FEATURE_CONTROL_LOCKED;
    test_bits |= FEATURE_CONTROL_VMXON_ENABLED_OUTSIDE_SMX;

    if ((old & test_bits) != test_bits) {
        wrmsrl(MSR_IA32_FEATURE_CONTROL, old | test_bits);
    }

    cr4_set_bits(X86_CR4_VMXE);
    return 0;
}

int vmxon() {
    struct vmcs *vmxon;
    vmxon = this_cpu_ptr(&vmxarea);
    uint32_t vmx_msr_low, vmx_msr_high;
    rdmsr(MSR_IA32_VMX_BASIC, &vmx_msr_low, &vmx_msr_high);
    revision_id = vmx_msr_low & 0x7fffffff;
    pr_info("revision_id: 0x%x\n", revision_id);
    vmxon->revision_id = revision_id;

    native_vmxon(__pa(vmxon));
    pr_info("vmxon succeed\n");
    return 0;
}

int setup_vmcs_control() {
    uint32_t _pin_based_exec_control = 0;

    adjust_vmx_controls(0, 0, MSR_IA32_VMX_PINBASED_CTLS,
                        &_pin_based_exec_control);
    pr_info("_pin_based_exec_control: 0x%x\n", _pin_based_exec_control);

    uint32_t _cpu_based_exec_control = 0;
    adjust_vmx_controls(0,
                        0,
                        MSR_IA32_VMX_PROCBASED_CTLS,
                        &_cpu_based_exec_control);

    uint32_t _vm_entry_control = 0;
    adjust_vmx_controls(VM_ENTRY_IA32E_MODE,
                       0, MSR_IA32_VMX_ENTRY_CTLS,
                        &_vm_entry_control);

    uint32_t _vm_exit_control = 0;
    adjust_vmx_controls(VM_EXIT_HOST_ADDR_SPACE_SIZE,
                        0,
                        MSR_IA32_VMX_EXIT_CTLS,
                        &_vm_exit_control);

    vmcs_write32(PIN_BASED_VM_EXEC_CONTROL, _pin_based_exec_control);
    vmcs_write32(CPU_BASED_VM_EXEC_CONTROL, _cpu_based_exec_control);
    vmcs_write32(VM_EXIT_CONTROLS, _vm_exit_control);
    vmcs_write32(VM_ENTRY_CONTROLS, _vm_entry_control);
    return 0;
}

void vm_exit_embedded() {
    pr_info("vm_exit_embedded say hi\n");

    pr_info("guest_rip: 0x%lx\n", vmcs_readl(GUEST_RIP));

    uint32_t exit_reason = vmcs_read32(VM_EXIT_REASON);
    pr_info("exit_reason: %d\n", exit_reason);
    //
    uint32_t entry_failure = (exit_reason >> 31) & 1;
    if (entry_failure) {
        pr_info("entry failure.\n");
    }
}

static void vmx_vcpu_run() {
    asm volatile(
            "1: "
            "mov %rax, 0\n\t"
            "mov %rdi, 2\n\t"
            "mov %rsi, 3\n\t"
            "mov %rdx, 3\n\t"
            "mov %rcx, 3\n\t"
            "vmcall\n\t"
            "jmp 1b\n\t"
            "2: "
            "call vm_exit_embedded\n\t"

            ".global vm_exit_func \n\t"
            "vm_exit_func: .quad 2b \n\t"

            ".global guest_entry \n\t"
            "guest_entry: .quad 1b \n\t"
            );
}

static inline unsigned long get_tr_base(void)
{
    struct desc_ptr dt;
    struct segment_desc *gdt;
    struct tss_desc *desc;

    store_gdt(&dt);
    gdt = (void *)dt.address;
    desc = (void *)&gdt[store_tr() / 8];
    BUG_ON(desc->p != 1);
    return ((unsigned long)desc->base0) |
           ((unsigned long)desc->base1 << 16) |
           ((unsigned long)desc->base2 << 24) |
           ((unsigned long)desc->base3 << 32);
}


int setup_vmcs_host_state() {
    vmcs_writel(HOST_CR0, read_cr0());
    vmcs_writel(HOST_CR3, read_cr3());
    vmcs_writel(HOST_CR4, read_cr4());

    vmcs_write16(HOST_CS_SELECTOR, read_cs());
    vmcs_write16(HOST_DS_SELECTOR, read_ds());
    vmcs_write16(HOST_ES_SELECTOR, read_es());
    vmcs_write16(HOST_SS_SELECTOR, read_ss());
    vmcs_write16(HOST_TR_SELECTOR, store_tr());
    pr_info("vm_exit_func: 0x%lx\n", vm_exit_func);
    vmcs_writel(HOST_RIP,  vm_exit_func);
    pr_info("host_stack_top: 0x%lx\n", (uint64_t) host_stack);
    vmcs_writel(HOST_RSP, (uint64_t) host_stack + 0x1000);
    return 0;
}

int setup_vmcs_guest_state() {
    vmcs_writel(GUEST_CR0, read_cr0());
    vmcs_writel(GUEST_CR3, read_cr3());
    vmcs_writel(GUEST_CR4, read_cr4());

//    vmcs_write16(GUEST_CS_SELECTOR, 0xf000);
//    vmcs_writel(GUEST_CS_BASE, 0xffff0000ul);
//    vmcs_write32(GUEST_CS_LIMIT, 0xffff);
    vmcs_write32(GUEST_CS_AR_BYTES, 0x209b);

//    vmcs_write16(GUEST_DS_SELECTOR, 0);
//    vmcs_writel(GUEST_DS_BASE, 0);
//    vmcs_write32(GUEST_DS_LIMIT, 0xffff);
    vmcs_write32(GUEST_DS_AR_BYTES, 0x93);

//    vmcs_write16(GUEST_ES_SELECTOR, 0);
//    vmcs_writel(GUEST_ES_BASE, 0);
//    vmcs_write32(GUEST_ES_LIMIT, 0xffff);
    vmcs_write32(GUEST_ES_AR_BYTES, 0x93);

//    vmcs_write16(GUEST_SS_SELECTOR, 0);
//    vmcs_writel(GUEST_SS_BASE, 0);
//    vmcs_write32(GUEST_SS_LIMIT, 0xffff);
    vmcs_write32(GUEST_SS_AR_BYTES, 0x93);

//    vmcs_write16(GUEST_FS_SELECTOR, 0);
//    vmcs_writel(GUEST_FS_BASE, 0);
//    vmcs_write32(GUEST_FS_LIMIT, 0xffff);
    vmcs_write32(GUEST_FS_AR_BYTES, 0x93);

//    vmcs_write16(GUEST_GS_SELECTOR, 0);
//    vmcs_writel(GUEST_GS_BASE, 0);
//    vmcs_write32(GUEST_GS_LIMIT, 0xffff);
    vmcs_write32(GUEST_GS_AR_BYTES, 0x93);

//    vmcs_write16(GUEST_TR_SELECTOR, 0);
//    vmcs_writel(GUEST_TR_BASE, 0);
//    vmcs_write32(GUEST_TR_LIMIT, 0xffff);
    vmcs_write32(GUEST_TR_AR_BYTES, 0x008b);

//    vmcs_write16(GUEST_LDTR_SELECTOR, 0);
    vmcs_writel(GUEST_LDTR_BASE, 0);
    vmcs_write32(GUEST_LDTR_LIMIT, 0xffff);
    vmcs_write32(GUEST_LDTR_AR_BYTES, 0x00082);

    vmcs_write32(GUEST_SYSENTER_CS, 0);
    vmcs_writel(GUEST_SYSENTER_ESP, 0);
    vmcs_writel(GUEST_SYSENTER_EIP, 0);
    vmcs_write64(GUEST_IA32_DEBUGCTL, 0);

    vmcs_writel(GUEST_RFLAGS, 0x02);
    pr_info("guest_entry: 0x%lx\n", guest_entry);
    vmcs_writel(GUEST_RIP, guest_entry);

    vmcs_writel(GUEST_GDTR_BASE, 0);
    vmcs_write32(GUEST_GDTR_LIMIT, 0xffff);

    vmcs_writel(GUEST_IDTR_BASE, 0);
    vmcs_write32(GUEST_IDTR_LIMIT, 0xffff);
    vmcs_write64(VMCS_LINK_POINTER, ~UINT64_C(0));

    return 0;
}

int setup_vmcs() {
    struct vmcs *vmcs;
    vmcs = this_cpu_ptr(&current_vmcs);
    vmcs->revision_id = revision_id;
    vmcs_clear(__pa(vmcs));
    vmcs_load(__pa(vmcs));

    setup_vmcs_control();
    setup_vmcs_host_state();
    setup_vmcs_guest_state();
    return 0;
}

void just_for_break() {

}


void hypervisor_init(void) {
    pr_info("hypervisor: hello\n");

    if (!cpu_has_vmx()) {
        pr_info("this cpu don't have vmx.\n");
    }
    pr_info("cpu has vmx feature support\n");
    vmx_enable();
    vmxon();

    setup_vmcs();
    just_for_break();
    asm volatile ("vmlaunch;");


    pr_info("err_code: %d\n", vmcs_read32(VM_INSTRUCTION_ERROR));
}
